from ..StrategiesConstructor import StrategiesConstructor
from geesibling.core.types import Graph, Node
from ..shard.placement_types import PlacementStrategy
from .cost_graph import CostGraph
from .solver import Solver
from typing import Any, cast, Dict, List, Optional, Sequence, Tuple
# TODO
def solve_solution(gm: Graph, strategy_constructor: StrategiesConstructor):
    '''
    This method is used to solve the best solution for the given graph.
    The solution is a list of integers, each integer represents the best strategy index of the corresponding node.
    '''
    # temporarily we use all nodes as liveness list, we count the backward memory cost together with
    # forward memory cost into the node memory cost, and no activation checkpoint is used in this phase.
    # graph_analyser = GraphAnalyser(gm)
    # liveness_list = graph_analyser.liveness_analysis()
    cost_graph = CostGraph(strategy_constructor.leaf_strategies,strategy_constructor.strategy_map)

    solver = Solver(gm, strategy_constructor, cost_graph)
    ret = solver.call_solver_serialized_args()
    solution = list(ret[0])
    # strategy_map = solution_annotation_pass(solution,strategy_constructor)
    return solution

def solution_annotation_pass(solution: List[int],
                             strategy_constructor: StrategiesConstructor)->Dict[str,PlacementStrategy]:
    # 解析得到的solution
    nodes = [strategies_vector.node for strategies_vector in strategy_constructor.leaf_strategies]
    strategy_map = {}   #node:best_strategy
    for node_index, (node, strategy_index) in enumerate(zip(nodes, solution)):
        strategies_vector = strategy_constructor.strategy_map[node.name]
        strategy_map[node.name] = strategies_vector[strategy_index].sharding_specs
    return strategy_map
